from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

NUM_CLASSES = 6 
NUM_FEATURES = 2
colors = ('green', 'orange', 'blue', 'magenta', 'cyan', 'darkgray')

fig,axs = plt.subplots(1,4,figsize=(22,4), 
                       subplot_kw={'xticks':(), 'yticks':()})

axs[0].scatter(X[:,0], X[:,1], c = y)
axs[0].set_title('samples')
for max_iter in [1,2,3]:
    X,y = make_blobs(n_samples=100,n_features=NUM_FEATURES,centers=NUM_CLASSES,
           cluster_std=1,random_state=10)
    
    kmeans_clustering = KMeans(n_clusters = NUM_CLASSES, 
                               max_iter = max_iter, init = 'random')
    clusters = kmeans_clustering.fit_predict(X)
    
    axs[max_iter].set_title('KMeans, max_iter='+str(max_iter))
    for i in range(len(clusters)):
        color = colors[clusters[i]]
        axs[max_iter].scatter(X[i,0],X[i,1], c = color) 
        
    cluster_centers = kmeans_clustering.cluster_centers_ 
    for i in range(len(cluster_centers)):
        axs[max_iter].scatter(cluster_centers[:,0], cluster_centers[:,1],
                              marker='*', c='r', s=200)  
